{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.set_option(\"max_colwidth\", None)\n",
    "pd.set_option(\"max_seq_items\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "from autoflow.knowledge_graph.programs.extract_graph import (\n",
    "    KnowledgeGraphExtractor,\n",
    "    KnowledgeGraphOutput,\n",
    ")\n",
    "from autoflow.models.llms import LLM\n",
    "from autoflow.models.llms.dspy import get_dspy_lm_by_llm\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "llm = LLM()\n",
    "dspy_lm = get_dspy_lm_by_llm(llm)\n",
    "extractor = KnowledgeGraphExtractor(dspy_lm=dspy_lm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "text = Path(\"../tests/fixtures/tidb-overview.md\").read_text()\n",
    "actual_knowledge = extractor.forward(text)\n",
    "df = actual_knowledge.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>description</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB is an open-source distributed SQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads. It is MySQL compatible and features horizontal scalability, strong consistency, and high availability. TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TiDB Self-Managed</td>\n",
       "      <td>TiDB Self-Managed is a product option of TiDB that allows users or organizations to deploy and manage TiDB on their own infrastructure with complete flexibility. Users can enjoy the power of open-source distributed SQL while retaining full control over their environment.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Hybrid Transactional and Analytical Processing (HTAP)</td>\n",
       "      <td>HTAP is a database processing technique that enables both transactional and analytical workloads to be performed on the same database system without the need for separate systems or data duplication.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>OLTP</td>\n",
       "      <td>OLTP (Online Transactional Processing) is a class of systems that facilitate and manage transaction-oriented applications, typically for data entry and retrieval transactions in a database.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>OLAP</td>\n",
       "      <td>OLAP (Online Analytical Processing) is a category of software tools that provide analysis of data for business decisions. It allows users to analyze multidimensional data interactively from multiple perspectives.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>TiKV</td>\n",
       "      <td>TiKV is a row-based storage engine provided by TiDB for storing data in a distributed manner.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>TiFlash</td>\n",
       "      <td>TiFlash is a columnar storage engine provided by TiDB for real-time replication of data from TiKV, ensuring consistent data between row-based and columnar storage.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>TiDB Operator</td>\n",
       "      <td>TiDB Operator is a tool that helps manage TiDB on Kubernetes and automates tasks related to operating the TiDB cluster, making TiDB easier to deploy on any cloud that provides managed Kubernetes.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>TiDB Cloud</td>\n",
       "      <td>TiDB Cloud is a fully-managed TiDB service that provides the easiest, most economical, and most resilient way to deploy and run TiDB clusters in the cloud.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>MySQL Protocol</td>\n",
       "      <td>The MySQL Protocol is the protocol used by MySQL database systems for communication between clients and servers. TiDB is compatible with the MySQL protocol, allowing seamless integration with the MySQL ecosystem.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    name  \\\n",
       "0                                                   TiDB   \n",
       "1                                      TiDB Self-Managed   \n",
       "2  Hybrid Transactional and Analytical Processing (HTAP)   \n",
       "3                                                   OLTP   \n",
       "4                                                   OLAP   \n",
       "5                                                   TiKV   \n",
       "6                                                TiFlash   \n",
       "7                                          TiDB Operator   \n",
       "8                                             TiDB Cloud   \n",
       "9                                         MySQL Protocol   \n",
       "\n",
       "                                                                                                                                                                                                                                                                                                                           description  \n",
       "0  TiDB is an open-source distributed SQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads. It is MySQL compatible and features horizontal scalability, strong consistency, and high availability. TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.  \n",
       "1                                                      TiDB Self-Managed is a product option of TiDB that allows users or organizations to deploy and manage TiDB on their own infrastructure with complete flexibility. Users can enjoy the power of open-source distributed SQL while retaining full control over their environment.  \n",
       "2                                                                                                                              HTAP is a database processing technique that enables both transactional and analytical workloads to be performed on the same database system without the need for separate systems or data duplication.  \n",
       "3                                                                                                                                        OLTP (Online Transactional Processing) is a class of systems that facilitate and manage transaction-oriented applications, typically for data entry and retrieval transactions in a database.  \n",
       "4                                                                                                                 OLAP (Online Analytical Processing) is a category of software tools that provide analysis of data for business decisions. It allows users to analyze multidimensional data interactively from multiple perspectives.  \n",
       "5                                                                                                                                                                                                                                        TiKV is a row-based storage engine provided by TiDB for storing data in a distributed manner.  \n",
       "6                                                                                                                                                                  TiFlash is a columnar storage engine provided by TiDB for real-time replication of data from TiKV, ensuring consistent data between row-based and columnar storage.  \n",
       "7                                                                                                                                  TiDB Operator is a tool that helps manage TiDB on Kubernetes and automates tasks related to operating the TiDB cluster, making TiDB easier to deploy on any cloud that provides managed Kubernetes.  \n",
       "8                                                                                                                                                                          TiDB Cloud is a fully-managed TiDB service that provides the easiest, most economical, and most resilient way to deploy and run TiDB clusters in the cloud.  \n",
       "9                                                                                                                 The MySQL Protocol is the protocol used by MySQL database systems for communication between clients and servers. TiDB is compatible with the MySQL protocol, allowing seamless integration with the MySQL ecosystem.  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"entities\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>source_entity</th>\n",
       "      <th>relationship_desc</th>\n",
       "      <th>target_entity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB Self-Managed is a product option of TiDB that allows users or organizations to deploy and manage TiDB on their own infrastructure with complete flexibility.</td>\n",
       "      <td>TiDB Self-Managed</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.</td>\n",
       "      <td>OLTP</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.</td>\n",
       "      <td>OLAP</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.</td>\n",
       "      <td>Hybrid Transactional and Analytical Processing (HTAP)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB provides two storage engines: TiKV, a row-based storage engine, and TiFlash, a columnar storage engine.</td>\n",
       "      <td>TiKV</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB provides two storage engines: TiKV, a row-based storage engine, and TiFlash, a columnar storage engine.</td>\n",
       "      <td>TiFlash</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>TiDB</td>\n",
       "      <td>TiDB is compatible with the MySQL protocol, allowing seamless integration with the MySQL ecosystem.</td>\n",
       "      <td>MySQL Protocol</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>TiDB Cloud</td>\n",
       "      <td>TiDB Cloud is a fully-managed TiDB service that provides the easiest, most economical, and most resilient way to deploy and run TiDB clusters in the cloud.</td>\n",
       "      <td>TiDB</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>TiDB Operator</td>\n",
       "      <td>TiDB Operator helps manage TiDB on Kubernetes and automates tasks related to operating the TiDB cluster, making TiDB easier to deploy on any cloud that provides managed Kubernetes.</td>\n",
       "      <td>TiDB</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   source_entity  \\\n",
       "0           TiDB   \n",
       "1           TiDB   \n",
       "2           TiDB   \n",
       "3           TiDB   \n",
       "4           TiDB   \n",
       "5           TiDB   \n",
       "6           TiDB   \n",
       "7     TiDB Cloud   \n",
       "8  TiDB Operator   \n",
       "\n",
       "                                                                                                                                                                      relationship_desc  \\\n",
       "0                     TiDB Self-Managed is a product option of TiDB that allows users or organizations to deploy and manage TiDB on their own infrastructure with complete flexibility.   \n",
       "1                                                                                         TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.   \n",
       "2                                                                                         TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.   \n",
       "3                                                                                         TiDB provides users with a one-stop database solution covering OLTP, OLAP, and HTAP services.   \n",
       "4                                                                          TiDB provides two storage engines: TiKV, a row-based storage engine, and TiFlash, a columnar storage engine.   \n",
       "5                                                                          TiDB provides two storage engines: TiKV, a row-based storage engine, and TiFlash, a columnar storage engine.   \n",
       "6                                                                                   TiDB is compatible with the MySQL protocol, allowing seamless integration with the MySQL ecosystem.   \n",
       "7                           TiDB Cloud is a fully-managed TiDB service that provides the easiest, most economical, and most resilient way to deploy and run TiDB clusters in the cloud.   \n",
       "8  TiDB Operator helps manage TiDB on Kubernetes and automates tasks related to operating the TiDB cluster, making TiDB easier to deploy on any cloud that provides managed Kubernetes.   \n",
       "\n",
       "                                           target_entity  \n",
       "0                                      TiDB Self-Managed  \n",
       "1                                                   OLTP  \n",
       "2                                                   OLAP  \n",
       "3  Hybrid Transactional and Analytical Processing (HTAP)  \n",
       "4                                                   TiKV  \n",
       "5                                                TiFlash  \n",
       "6                                         MySQL Protocol  \n",
       "7                                                   TiDB  \n",
       "8                                                   TiDB  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[\"relationships\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_knowledge = KnowledgeGraphOutput(\n",
    "    entities=[\n",
    "        {\"name\": \"TiDB\", \"description\": \"A distributed SQL database\"},\n",
    "        {\"name\": \"TiKV\", \"description\": \"TiKV is a row-based storage engine.\"},\n",
    "        {\"name\": \"TiFlash\", \"description\": \"TiFlash is a column-based storage engine.\"},\n",
    "        {\n",
    "            \"name\": \"OLTP\",\n",
    "            \"description\": \"OLTP is a type of database that is optimized for transaction processing.\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"OLAP\",\n",
    "            \"description\": \"OLAP is a type of database that is optimized for analytics.\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"HTAP\",\n",
    "            \"description\": \"HTAP is a type of database that is optimized for both transaction processing and analytics.\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"TiDB Self-Managed\",\n",
    "            \"description\": \"TiDB Self-Managed is a product option of TiDB.\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"TiDB Cloud\",\n",
    "            \"description\": \"TiDB Cloud is a fully-managed TiDB service.\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"TiDB Operator\",\n",
    "            \"description\": \"TiDB Operator is a tool that helps manage TiDB on Kubernetes.\",\n",
    "        },\n",
    "    ],\n",
    "    relationships=[\n",
    "        {\n",
    "            \"source_entity\": \"TiDB\",\n",
    "            \"target_entity\": \"TiKV\",\n",
    "            \"relationship_desc\": \"TiDB uses TiKV as its storage engine\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB\",\n",
    "            \"target_entity\": \"TiFlash\",\n",
    "            \"relationship_desc\": \"TiDB uses TiFlash as its analytics engine\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB\",\n",
    "            \"target_entity\": \"OLTP\",\n",
    "            \"relationship_desc\": \"TiDB supports OLTP workloads\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB\",\n",
    "            \"target_entity\": \"OLAP\",\n",
    "            \"relationship_desc\": \"TiDB supports OLAP workloads\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB\",\n",
    "            \"target_entity\": \"HTAP\",\n",
    "            \"relationship_desc\": \"TiDB supports HTAP workloads\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB Self-Managed\",\n",
    "            \"target_entity\": \"TiDB\",\n",
    "            \"relationship_desc\": \"TiDB Self-Managed is a product option of TiDB\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB Cloud\",\n",
    "            \"target_entity\": \"TiDB\",\n",
    "            \"relationship_desc\": \"TiDB Cloud is a fully-managed TiDB service.\",\n",
    "        },\n",
    "        {\n",
    "            \"source_entity\": \"TiDB Operator\",\n",
    "            \"target_entity\": \"TiDB Cloud\",\n",
    "            \"relationship_desc\": \"TiDB Operator is a tool that helps manage TiDB on Kubernetes.\",\n",
    "        },\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-11T09:00:13.095249Z",
     "start_time": "2025-04-11T09:00:11.025732Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6111111111111112"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from autoflow.knowledge_graph.programs.eval_graph import KnowledgeGraphEvaluator\n",
    "\n",
    "evaluator = KnowledgeGraphEvaluator(dspy_lm=dspy_lm)\n",
    "\n",
    "evaluation_result = evaluator.forward(expected_knowledge, actual_knowledge)\n",
    "evaluation_result.score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "\"EntityOutput\" object has no field \"metadata\"",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mValueError\u001b[39m                                Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[9]\u001b[39m\u001b[32m, line 4\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mautoflow\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mknowledge_graph\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mprograms\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mextract_covariates\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m EntityCovariateExtractor\n\u001b[32m      3\u001b[39m metadata_extractor = EntityCovariateExtractor(dspy_lm=dspy_lm)\n\u001b[32m----> \u001b[39m\u001b[32m4\u001b[39m actual_knowledge.entities = \u001b[43mmetadata_extractor\u001b[49m\u001b[43m.\u001b[49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtext\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mactual_knowledge\u001b[49m\u001b[43m.\u001b[49m\u001b[43mentities\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      6\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m entity \u001b[38;5;129;01min\u001b[39;00m actual_knowledge.entities:\n\u001b[32m      7\u001b[39m     \u001b[38;5;28mprint\u001b[39m(entity.name, entity.metadata)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/autoflow.ai/core/autoflow/knowledge_graph/programs/extract_covariates.py:81\u001b[39m, in \u001b[36mEntityCovariateExtractor.forward\u001b[39m\u001b[34m(self, text, entities)\u001b[39m\n\u001b[32m     78\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m entity \u001b[38;5;129;01min\u001b[39;00m entities:\n\u001b[32m     79\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m entity.name \u001b[38;5;129;01min\u001b[39;00m output_entity_map:\n\u001b[32m     80\u001b[39m         \u001b[38;5;66;03m# Update the covariates in the metadata of the entity.\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m81\u001b[39m         \u001b[43mentity\u001b[49m\u001b[43m.\u001b[49m\u001b[43mmetadata\u001b[49m = output_entity_map[entity.name].covariates\n\u001b[32m     83\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m entities\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/autoflow.ai/core/.venv/lib/python3.12/site-packages/pydantic/main.py:925\u001b[39m, in \u001b[36mBaseModel.__setattr__\u001b[39m\u001b[34m(self, name, value)\u001b[39m\n\u001b[32m    922\u001b[39m     \u001b[38;5;28mself\u001b[39m.__pydantic_validator__.validate_assignment(\u001b[38;5;28mself\u001b[39m, name, value)\n\u001b[32m    923\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.model_config.get(\u001b[33m'\u001b[39m\u001b[33mextra\u001b[39m\u001b[33m'\u001b[39m) != \u001b[33m'\u001b[39m\u001b[33mallow\u001b[39m\u001b[33m'\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m name \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.__pydantic_fields__:\n\u001b[32m    924\u001b[39m     \u001b[38;5;66;03m# TODO - matching error\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m925\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[33mf\u001b[39m\u001b[33m'\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m.\u001b[34m__class__\u001b[39m.\u001b[34m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m\u001b[33m object has no field \u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m\u001b[33m'\u001b[39m)\n\u001b[32m    926\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.model_config.get(\u001b[33m'\u001b[39m\u001b[33mextra\u001b[39m\u001b[33m'\u001b[39m) == \u001b[33m'\u001b[39m\u001b[33mallow\u001b[39m\u001b[33m'\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m name \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.__pydantic_fields__:\n\u001b[32m    927\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.model_extra \u001b[38;5;129;01mand\u001b[39;00m name \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.model_extra:\n",
      "\u001b[31mValueError\u001b[39m: \"EntityOutput\" object has no field \"metadata\""
     ]
    }
   ],
   "source": [
    "from autoflow.knowledge_graph.programs.extract_covariates import (\n",
    "    EntityCovariateExtractor,\n",
    ")\n",
    "\n",
    "metadata_extractor = EntityCovariateExtractor(dspy_lm=dspy_lm)\n",
    "actual_knowledge.entities = metadata_extractor.forward(text, actual_knowledge.entities)\n",
    "\n",
    "for entity in actual_knowledge.entities:\n",
    "    print(entity.name, entity.metadata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
